/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.ast;

import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import alloy.type.BasicType;
import alloy.type.RelationType;
import alloy.type.UnionType;
import alloy.util.Dbg;

/**
 * A SigmaExpr is an intermediate form of a n-arity union used by atomization
 * whose goal is to create generate SigmaExpr's from supertypes and union types
 * and then get rid of them at the level of formulas, decls, and intexpressions.
 * 
 * Each SigmaExpr categorizes the expressions it contains by atomic type.  A map
 * is created so that given a type, the "term" of the union of that type can be
 * easily accessed. 
 */
public class SigmaExpr extends TypedExpr {

    /** Maps a RelationType to the Expr of that type */
    private Map _typeToExpr;    
 
    /**
     * Creates a new SigmaExpr and adds the terms.
     * 
     * @param terms A list of <code>Node</code>s.  It is assumed that each expression  
     * is of an atomic type (or a single relation type that consists only of atomic types).
     * Unless otherwise stated, a "type" in this file refers to a RelationType.    
     */
    public SigmaExpr(List terms) {
        this((Node)terms.get(0));

        Iterator termsIter = terms.iterator();
        termsIter.next();
        while (termsIter.hasNext()) {
            addTerm((Node)termsIter.next());
        }
        
        //checkRep();
    }

    /**
     * Creates a new SigmaExpr
     * 
     * @param n1_ the first expression in the SigmaExpr.  Node should have a single atomic type.
     */
    public SigmaExpr(Node n1_) {
        super(n1_);
                
        Dbg.chk(n1_, "tried to create a SigmaExpr with a null node");
        if (n1_ instanceof TypedExpr) {
            UnionType type = ((TypedExpr)n1_).getType();
            Dbg.chk(
                type.size() == 1,
                "attempted to give an ill-typed node to SigmaExpr: " + n1_ + " " + n1_.getClass());
            _typeToExpr = new LinkedHashMap();
            _typeToExpr.put((RelationType)type.getTypeIter().next(), n1_);
            Set endTypeSet = new HashSet();
            endTypeSet.add(n1_);
            Set startTypeSet = new HashSet(endTypeSet);
            setType(type);
        } else {
            Dbg.fail("attempted to give a SigmaExpr an untyped child: " + n1_);
        }

        //checkRep();
    }

    /**
     * Adds a new term to this SigmaExpr.  If a term of the same type already exists, takes 
     * the union of those two terms.
     * 
     * @param n1_ the term to be added 
     * 
     */
    public void addTerm(Node n1_) {
        //checkRep();
        if (n1_ instanceof TypedExpr) {
            UnionType type = ((TypedExpr)n1_).getType();
            Dbg.chk(
                type.size() == 1 && type.arity() == ((TypedExpr)childAt(0)).getType().arity(),
                "attempted to add an ill-typed node to SigmaExpr: " + n1_);
            setType(getType().union(type));
            RelationType relType = (RelationType)type.getTypeIter().next();
            BasicType startType = relType.getBasicTypeAt(0);
            BasicType endType = relType.getBasicTypeAt(relType.arity() - 1);
            if (_typeToExpr.containsKey(relType)) {
                Expr expr = (Expr)_typeToExpr.get(relType);
                // remove old expr from startType and endType mappings, since
                // we're making a new object
                expr = expr.union((TypedExpr)n1_);
                _typeToExpr.put(relType, expr);
            } else {
                _typeToExpr.put(relType, n1_);
            }
        } else {
            Dbg.fail("attempted to give a SigmaExpr an untyped child: " + n1_);
        }
        addChild(n1_);
        //checkRep();
    }

    /**     
     * @return a Set of all the types for which a term has been stored.
     */
    public Set getTypes() {
        return _typeToExpr.keySet();
    }

    /**
     * 
     * @param type a RelationType.  It should consist of only atomic types; otherwise,
     * the method will always return false.
     * @return true if a term of type <code>type</code> exists in this SigmaExpr.
     */
    public boolean containsType(RelationType type) {
        return _typeToExpr.containsKey(type);
    }

    /**
     * 
     * @return the Set of atomic types (represented as SigType's) that appear on the 
     * first column of all RelationType's that are contained in this SigmaExpr. 
     */
    public Set getFirstColumnTypes() {
        Set firstColTypes = new HashSet(_typeToExpr.size());
        for (Iterator typeIter = _typeToExpr.keySet().iterator(); typeIter.hasNext();) {
            firstColTypes.add(((RelationType)typeIter.next()).getBasicTypeAt(0));
        }
        return firstColTypes;
    }

    /**
     * 
     * @param type a BasicType.  If not atomic, method will return false.
     * @return true if <code>type</code> appears in the first column of any types contained
     * in this SigmaExpr.
     */
    public boolean containsFirstColumnType(BasicType type) {
        for (Iterator typeIter = _typeToExpr.keySet().iterator(); typeIter.hasNext();) {
            RelationType relType = (RelationType)typeIter.next();
            if (relType.getBasicTypeAt(0)==type) {
                return true;
            }
        }
        return false;
    }

    /**
    * 
    * @return the Set of atomic types (represented as SigType's) that appear on the 
    * last column of all RelationType's that are contained in this SigmaExpr. 
    */
    public Set getLastColumnTypes() {
        Set lastColTypes = new HashSet(_typeToExpr.size());
        for (Iterator typeIter = _typeToExpr.keySet().iterator(); typeIter.hasNext();) {
            RelationType type = (RelationType)typeIter.next();
            lastColTypes.add(type.getBasicTypeAt(type.arity()-1));
        }
        return lastColTypes;
    }
    
    /**
    * 
    * @param type a BasicType.  If not atomic, method will return false.
    * @return true if <code>type</code> appears in the last column of any types contained
    * in this SigmaExpr.
    */    
    public boolean containsLastColumnType(BasicType type) {
        for (Iterator typeIter = _typeToExpr.keySet().iterator(); typeIter.hasNext();) {
            RelationType relType = (RelationType)typeIter.next();
            if (relType.getBasicTypeAt(relType.arity()-1)==type) {
                return true;
            }
        }
        return false;
    }

    /**
     * Calls getTermsWithFirstColumn(type, true);
     * 
     * @param type an atomic BasicType (returns an empty set, if not atomic).
     * @return a Set of Expr's whose type has <code>type</code> as its first column.  returns
     * an empty set if no such Expr exists.  Note that the returned Exprs are <b>copies</b>
     * of the original, and their LeafIds are not necessarily set.
     *      
     */        
    public Set getTermsWithFirstColumn(BasicType type) {
        return getTermsWithFirstColumn(type, true);
    }

    /**
     * @param type an atomic BasicType (returns an empty set, if not atomic).
     * @param copyTerms true to return copies of the terms
     * @return a Set of Expr's whose type has <code>type</code> as its first column.  returns
     * an empty set if no such Expr exists.  Returned expressions are the original and should
     * not be modified.  The method is used for easier reference equality comparisons.
     */
    public Set getTermsWithFirstColumn(BasicType type, boolean copyTerms) {
        Set termsSet = new HashSet();
        for (Iterator entryIter = _typeToExpr.entrySet().iterator(); entryIter.hasNext();) {
            Map.Entry entry = (Map.Entry)entryIter.next();
            RelationType relType = (RelationType)entry.getKey();
            BasicType curBasType = relType.getBasicTypeAt(0);
            Expr curTerm = (Expr)entry.getValue();
            if (curBasType == type) {
                termsSet.add(copyTerms? copyTerm(curTerm) : curTerm);                                    
            }
        }
        return termsSet;
    }

    /**
     * Calls getTermsWithLastColumn(type, true);
     * 
     * @param type an atomic BasicType (returns an empty set, if not atomic).
     * @return a Set of Expr's whose type has <code>type</code> as its last column.  returns
     * an empty set if no such Expr exists.  Note that the returned Exprs are <b>copies</b>
     * of the original, and their LeafIds are not necessarily set.
     *      
     */        
    public Set getTermsWithLastColumn(BasicType type) {
        return getTermsWithLastColumn(type, true);
    }
    
    /**
    * @param type an atomic BasicType (returns an empty set, if not atomic).
    * @param copyTerms true to return copies of the terms
    * @return a Set of Expr's whose type has <code>type</code> as its last column.  returns
    * an empty set if no such Expr exists.  Returned expressions are the original and should
    * not be modified.  The method is used for easier reference equality comparisons.
    */    
    public Set getTermsWithLastColumn(BasicType type, boolean copyTerms) {
        Set termsSet = new HashSet();
        for (Iterator entryIter = _typeToExpr.entrySet().iterator(); entryIter.hasNext();) {
            Map.Entry entry = (Map.Entry)entryIter.next();
            RelationType relType = (RelationType)entry.getKey();
            BasicType curBasType = relType.getBasicTypeAt(relType.arity()-1);
            Expr curTerm = (Expr)entry.getValue();
            if (curBasType == type) {
                termsSet.add(copyTerms? copyTerm(curTerm) : curTerm);                                    
            }
        }
        return termsSet;
    }
    
    /**
     * Calls getTerm(type, true)
     */
    public Expr getTerm(RelationType type) {
        return getTerm(type, true);
    }

    /**
     * 
     * @param type the atomic type to be looked up
     * @param copyTerm true if a copy should be returned
     * @return an Expr corresponding to the term of the desired type.  A copy is returned if 
     * copyTerm is set to true.
     */
    public Expr getTerm(RelationType type, boolean copyTerm) {
        Expr expr = (Expr)_typeToExpr.get(type);        
        return (copyTerm? copyTerm(expr) : expr);
    }
    
    /**
     * calls getTerms(true);
     */
    public Collection getTerms() {
        return getTerms(true);
    }
    
    /**
     * @param copyTerms true if copies should be returned
     * @return a Collection of all terms (as Expr's) stored in this SigmaExpr. 
     */
    public Collection getTerms(boolean copyTerms) {
        Collection c = new HashSet();
        for (Iterator iter = _typeToExpr.values().iterator(); iter.hasNext();) {
            Expr curTerm = (Expr)iter.next();
		    c.add(copyTerms? copyTerm(curTerm) : curTerm);
            
        }
        return c;
    }
    
    /**
     * you would want to use this when you are stuck with a SigmaExpr as your highest
     * level variable (e.g. this happens sometimes when using the evaluator) and you 
     * need to iterate through the terms, and the leafIdMaps should be non-null. 
     */
    public Iterator getUncopiedTermsIter() {
    	return _typeToExpr.values().iterator();
    }

    private Expr copyTerm(Expr term) {
        if (term == null) return null;
    	Expr copy = (Expr)term.copy();
        
        /*
         this code is no longer needed bcause during atomization FixLeafIdsVisitor will
         fill in null LeafIdMap's gracefully         
        if (term instanceof VarCreator) {
            ((VarCreator)copy).setLeafIdMap(((VarCreator)term).getLeafIdMap());
        }*/
        return copy;
    }       

    /**
     *  Returns a set containing the types that are contained in both 
     * this.getTypes() and sigExpr_.getTypes()
     */
    public Set getIntersectedTypes(SigmaExpr sigExpr_) {
        Set ret = new HashSet(getTypes());
        ret.retainAll(sigExpr_.getTypes());
        return ret;
    }

    public String nodeString() {
        return super.childrenStr("S[( ", " )+( ", " )]");
    }

    public String toString() {
        return nodeString();
    }

    /**
     * accept an AST visitor
     * @param visitor_ the visitor
     */
    protected void acceptVisitor(ASTVisitor visitor_) {
        visitor_.visit(this);
	/** EMINA:  THIS CAUSED BUGS IN THE EVALUATOR.  What's its purpose? **/
	/**
        for (Iterator valIter = _typeToExpr.values().iterator(); valIter.hasNext(); ) {
        	Expr expr = (Expr)valIter.next();
        	expr.applyVisitor(visitor_);
        }
	**/
    }

    /**
     * accept a returning AST visitor
     * @param visitor_ the visitor
     * @return the Object returned by the visitor
     */
    protected Object acceptReturnVisitor(ASTReturnVisitor visitor_) {
        return visitor_.visit(this);        
    }
    
    private boolean checkRep() {
        for (Iterator entryIter = _typeToExpr.entrySet().iterator(); entryIter.hasNext();) {
            Map.Entry entry = (Map.Entry) entryIter.next();
            RelationType type = (RelationType) entry.getKey();
            TypedExpr expr = (TypedExpr) entry.getValue();
            Dbg.chk(type.equals(expr.getType().toRelationType()));
        }
    
        return true;
    }

}
